import torch
from PIL import Image
from torchvision.utils import make_grid
import numpy as np
import matplotlib.pyplot as plot

def save_image_gray(tensor, filename, nrow=8, padding=2,
               normalize=False, range=None, scale_each=False, pad_value=0):
    """Save a given Tensor into an image file.

    Args:
        tensor (Tensor or list): Image to be saved. If given a mini-batch tensor,
            saves the tensor as a grid of images by calling ``make_grid``.
        **kwargs: Other arguments are documented in ``make_grid``.
    """
    grid = make_grid(tensor, nrow=nrow, padding=padding, pad_value=pad_value,
                     normalize=normalize, range=range, scale_each=scale_each)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid.mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
    im = Image.fromarray(ndarr)
    im.convert('L').save(filename)

def Normalize(x):
    x = np.array(x).astype(np.float32)
    x /=255.0
    x-=(0.485,0.456,0.406)
    x /=(0.229,0.224,0.225)
    return x

pic_path = './n0012.jpg'
pic = Image.open(pic_path).convert('RGB')

pic = pic.resize((512,512))
print('pic shape:{}'.format(pic.size))

pic = np.array(pic)
pic = Normalize(pic)

pic = np.transpose(pic,(2,0,1))
pic = torch.from_numpy(pic.copy()).float()
pic = pic.unsqueeze(0)


device = torch.device("cuda:0")
pic = pic.to(device)


danet = torch.load('./danet.pth')
danet = danet.to(device)
danet = danet.eval()


out = danet(pic)
out_all = out[0]
out_p = out[1]
out_c = out[2]

out = out_all

print(out.shape)
with torch.no_grad():
    save_image_gray(out, r'./test.png')

plot.figure()
res = np.array(Image.open(r'./test.png'))

plot.imshow(res)
plot.show()